import importlib
import inspect
import os
import re
import subprocess
from functools import partial
from datetime import datetime
import paramiko
import yaml
from colorama import Fore, Style, init

from plugins.common import Common, CheckerBase


class RemoteTest:

    def __init__(self):
        self.output_data = ""
        self.output_file = self.generate_logs()
        init()
        self.colors = {
            "black": Fore.BLACK,
            "red": Fore.RED,
            "green": Fore.GREEN,
            "yellow": Fore.YELLOW,
            "blue": Fore.BLUE,
            "magenta": Fore.MAGENTA,
            "cyan": Fore.CYAN,
            "white": Fore.WHITE,
            "reset": Style.RESET_ALL
        }
        self.plugin_classes = self.get_all_plugin_classes()

    @staticmethod
    def generate_logs(directory="logs"):
        current_datetime = datetime.now().strftime("%Y-%m-%d__%H.%M.%S")
        output_file = os.path.join(directory, f"{current_datetime}.log")
        return output_file

    def color_text(self, text, color):
        return self.colors[color] + text + self.colors["reset"]

    @staticmethod
    def read_yaml(file_path):
        file_path = f"tasks/{file_path}"
        with open(file_path, "r", encoding="utf-8") as f:
            return yaml.safe_load(f)

    @staticmethod
    def analyse_host(host_info):
        pattern = r'(?:(?P<username>[^:]+):(?P<password>.*?(?=##))##)?(?P<ip>[^:]+):(?P<port>\d+)(?:###(?P<remark>.*))?'
        match = re.match(pattern, host_info)
        if match:
            return match.group('username') or None, \
                   match.group('password') or None, \
                   match.group('ip'), \
                   match.group('port'), \
                   match.group('remark') or None
        else:
            return None, None, None, None, None

    @staticmethod
    def analyse_cmd(cmd_str):
        pattern = r'(?:(?:fun##)(?P<func_name>\w+)(?:\((?P<params>[^)]*)\))?)|(?P<shell_cmd>[^f][^u][^n].*)'
        match = re.match(pattern, cmd_str)
        if match:
            func_name = match.group('func_name')
            params = match.group('params')
            shell_cmd = match.group('shell_cmd')
            params_dict = {}
            if func_name:
                if params:
                    # params_dict = dict(param.split('=') for param in params.split(','))
                    for param in params.split(','):
                        parts = param.split('=')
                        if len(parts) != 2:
                            raise ValueError(f"Invalid parameter format: {param}")
                        key, value = parts
                        params_dict[key] = value
                # else:
                #     params_dict = {}
                return None, func_name, params_dict
            else:
                return shell_cmd, None, {}
        else:
            return None, None, {}

    @staticmethod
    def analyse_class(class_str):
        pattern = r'(?P<class_name>\w+)\((?P<class_params>.*)\)'
        match = re.match(pattern, class_str)

        if match:
            class_name = match.group('class_name')
            params_string = match.group('class_params')

            if params_string:
                params_pattern = r'(\w+)=(.*?)(?:,|$)'
                params_matches = re.findall(params_pattern, params_string)
                params_dict = {param: value.strip('\'"') for param, value in params_matches}
            else:
                params_dict = {}
        else:
            class_name = class_str
            params_dict = {}

        return class_name, params_dict

    def ping(self, ip):
        try:
            output = subprocess.check_output(["ping", "-n", "1", ip])
            if "TTL=" in output.decode("gbk"):
                return True
            else:
                print(self.color_text(" ping不通", "red"))
                self.output_data += " ping不通\n"
                return False
        except subprocess.CalledProcessError:
            print(self.color_text(" ping不通", "red"))
            self.output_data += " ping不通\n"
            return False

    def get_class(self, class_name):
        for cls in self.plugin_classes:
            if cls.__name__ == class_name:
                return cls
        return None

    @staticmethod
    def get_all_plugin_classes():
        plugin_modules = [
            importlib.import_module(f'plugins.{module_name[:-3]}')
            for module_name in os.listdir('plugins')
            if module_name.endswith('.py') and module_name != '__init__.py'
        ]
        plugin_classes = []
        for module in plugin_modules:
            for name, obj in inspect.getmembers(module):
                if inspect.isclass(obj):
                    plugin_classes.append(obj)
        return plugin_classes

    def run(self, private_key_path, num_threads=3, uuid=None):
        main_yaml = self.read_yaml("main.yml")

        for task in main_yaml["tasks"]:
            sub_yaml = self.read_yaml(task["include"])

            for host_info in sub_yaml[0]["hosts"]:
                ssh = paramiko.SSHClient()
                ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                user, password, ip, port, remark = self.analyse_host(host_info)
                print(self.color_text(f"{remark or ip}:", "green"), end="")
                self.output_data += f"\n{remark or ip}:"
                if not self.ping(ip):
                    continue

                try:
                    if user is not None and password is not None:
                        ssh.connect(ip, port=port, username=user, password=password)
                    else:
                        private_key = paramiko.RSAKey.from_private_key_file(private_key_path)
                        ssh.connect(ip, username="root", pkey=private_key, port=port)
                except (paramiko.ssh_exception.NoValidConnectionsError,
                        paramiko.SSHException,
                        Exception) as e:
                    print(self.color_text(f" ssh连接失败: {e}", "red"))
                    self.output_data += f" ssh连接失败: {e}\n"
                    continue

                print("")
                self.output_data += "\n"
                partial_funcs = []
                all_commands_info = []

                for command_info in sub_yaml[0]["commands"]:
                    if uuid and uuid != command_info.get('uuid', None):
                        continue

                    all_commands_info.append(command_info)
                    class_str = command_info.get('class', None)

                    if class_str:
                        class_name, class_params_dict = self.analyse_class(class_str)
                        class_params_dict['ssh_client'] = ssh
                        cls = self.get_class(class_name)
                        if cls:
                            checker = cls(**class_params_dict)
                        else:
                            raise RuntimeError(f"未找到类{class_name}")
                    else:
                        checker = Common(ssh_client=ssh)

                    cmd_str = command_info["cmd"]
                    cmd, func_name, params_dict = self.analyse_cmd(cmd_str)

                    if func_name:
                        func = getattr(checker, func_name)
                        partial_fun = partial(func, **params_dict)
                    else:
                        func = getattr(checker, "run_command")
                        partial_fun = partial(func, cmd)

                    partial_funcs.append(partial_fun)

                try:
                    results = CheckerBase.run_checks(partial_funcs, num_threads)

                    for (result, result_err), command_info in zip(results, all_commands_info):
                        cmd_name = self.color_text(command_info["name"] + ":", "yellow")
                        print("  {:<40}".format(cmd_name), end="", flush=True)
                        self.output_data += "  {:<40}".format(command_info["name"] + ":")
                        passed, expect_str = "", ""
                        expect_value = command_info.get("expect_value", None)
                        expect_equation = command_info["expect_equation"]

                        if expect_equation == "nc":
                            expect_str_color = self.color_text(f"{result_err}", "red")
                            out = result if result else expect_str_color
                            print("执行结果: {:<13}".format(out.replace("\n", " ")), self.color_text("[不校验]", "blue"))
                            self.output_data += "执行结果: {:<13} [不校验]\n".format(out.replace("\n", " "))
                            continue

                        if expect_equation == "eq":
                            passed = int(result == str(expect_value))
                            expect_str = "(eq {})".format(expect_value)
                        elif expect_equation == "ge":
                            passed = int(float(result) >= float(expect_value))
                            expect_str = "(ge {})".format(expect_value)
                        elif expect_equation == "le":
                            passed = int(float(result) <= float(expect_value))
                            expect_str = "(le {})".format(expect_value)
                        elif expect_equation == "re":
                            passed = int(re.match(expect_value, result) is not None)
                            expect_str = "(re {})".format(expect_value)

                        print("预期值: {:<8} 实际值: {:<5} ".format(expect_str, result), end="", flush=True)
                        self.output_data += "预期值: {:<8} 实际值: {:<5} ".format(expect_str, result)
                        print(self.color_text("[通过]" if passed else "[不通过]", "green" if passed else "red"))
                        self.output_data += "[通过]\n" if passed else "[不通过]\n"

                except Exception as e:
                    print("执行异常: {}".format(e), end="")
                    self.output_data += "执行异常: {}".format(e)
                    print(self.color_text("[不通过]", "red"))
                    self.output_data += "[不通过]"

                ssh.close()
        with open(self.output_file, "w", encoding="utf-8") as f:
            f.write(self.output_data)


if __name__ == "__main__":
    runner = RemoteTest()
    runner.run(num_threads=5, private_key_path=r'C:\Users\love\.ssh\id_rsa')
